# -*- coding: utf-8 -*-
"""
21 维 LHS 采样
变量:
  FLR01_DUR, FLR02_DUR, FLR18_DUR   : 非标准层  PERT(5,6,9)
  FLR03_DUR ... FLR17_DUR           : 标准层    PERT(4,5,8)  共15层
  DLV_DLY                            : 截断正态
  CRN_FAL                            : 泊松故障序列
  RAIN_HOLD, WIND_HOLD               : 逐日伯努利矩阵
"""
import pandas as pd
import numpy as np
from scipy.stats import qmc, truncnorm, beta, poisson

# ---------- 1. 读历史频率 ----------
def load_daily_prob(csv_file):
    df = pd.read_csv(csv_file, usecols=['Month', 'Day', 'frequency'])
    df.columns = ['Month', 'Day', 'freq']
    return {(int(r.Month), int(r.Day)): r.freq for _, r in df.iterrows()}

# 加载降水和大风概率
rain_prob = load_daily_prob('Precipiation_wind_stat.csv')
wind_prob = load_daily_prob('strong_wind_stat.csv')

# ---------- 2. PERT 辅助 ----------
def pert_rvs(a, b, c, size):
    alpha = (4 * b + c - 5 * a) / (c - a)
    beta_param = (5 * c - 4 * b - a) / (c - a)
    return a + (c - a) * beta.rvs(alpha, beta_param, size=size)

# ---------- 3. LHS ----------
def lhs_sample(n=1000):
    dim = 21
    sampler = qmc.LatinHypercube(d=dim, seed=42)
    u = sampler.random(n)

    # 非标准层 3 个
    flr01 = pert_rvs(5, 6, 9, size=n)
    flr02 = pert_rvs(5, 6, 9, size=n)
    flr18 = pert_rvs(5, 6, 9, size=n)

    # 标准层 15 个
    flr03_17 = np.column_stack([pert_rvs(4, 5, 8, size=n) for _ in range(15)])

    # 构件延迟
    a, b = (-10 + 4) / 2, (10 + 4) / 2
    dlv_dly = truncnorm.ppf(u[:, 18], a, b, loc=-4, scale=2)

    # 塔吊故障序列
    T_hours = 365 * 10  # 假设每天工作 10 小时
    lam = 1 / 3000
    E_N = lam * T_hours
    N_low = max(1, int(E_N - 2 * np.sqrt(E_N)))
    N_high = int(E_N + 2 * np.sqrt(E_N))
    N = np.round(N_low + u[:, 19] * (N_high - N_low)).astype(int)
    crn_fal = [np.sort(np.random.rand(k) * T_hours).round(2).tolist() for k in N]  # 保留两位小数

    # 逐日停工事件
    def daily_events(prob_dict):
        events = np.empty((n, 365), dtype=bool)
        for day_idx, d in enumerate(pd.date_range('2025-03-01', periods=365, freq='D')):
            p = prob_dict.get((d.month, d.day), 0.0)
            s = qmc.LatinHypercube(d=1, seed=2000 + day_idx)
            events[:, day_idx] = s.random(n)[:, 0] < p
        return events

    rain_hold = daily_events(rain_prob)
    wind_hold = daily_events(wind_prob)

    # 构造 DataFrame
    df = pd.DataFrame({
        'FLR01_DUR': flr01,
        'FLR02_DUR': flr02,
        'FLR18_DUR': flr18,
        **{f'FLR{i:02d}_DUR': flr03_17[:, i - 3] for i in range(3, 18)},
        'DLV_DLY': dlv_dly,
        'CRN_FAL': crn_fal,
        'RAIN_HOLD': list(rain_hold),
        'WIND_HOLD': list(wind_hold)
    })

    # 保留两位小数
    df = df.round({'FLR01_DUR': 2, 'FLR02_DUR': 2, 'FLR18_DUR': 2, **{f'FLR{i:02d}_DUR': 2 for i in range(3, 18)}, 'DLV_DLY': 2})

    return df

# ---------- 4. 运行 ----------
if __name__ == '__main__':
    df = lhs_sample(1000)
    print(df.head())
    df.to_csv('lhs_21vars.csv', index=False, float_format='%.2f')  # 保存时指定格式化方式
    df.to_pickle('lhs_21vars.pkl')